import torch.nn as nn
from .output import *
from .decoder import *
# from .input import *
# from .encoder import *
# from .element import *


class MyTransformer(nn.Module):
    def __init__(self, en_embed, encoder, de_embed, decoder, generator):
        super().__init__()
        self.en_embed = en_embed
        self.encoder = encoder
        self.de_embed = de_embed
        self.decoder = decoder
        self.generator = generator

    def forward(self, source_x, target_y, source_mask, target_mask):
        en_position_data = self.en_embed(source_x)
        en_data = self.encoder(en_position_data, source_mask)
        de_position_data = self.de_embed(target_y)
        de_data = self.decoder(
            data=de_position_data,
            encoder_output=en_data,
            source_mask=source_mask,
            target_mask=target_mask,
        )
        return self.generator(de_data)
